from transformers import AutoProcessor, AutoModel
import torch
from PIL import Image

class ImageProcessor:
    def __init__(self):
        self.processor = AutoProcessor.from_pretrained("OFA-Sys/chinese-clip-vit-base-patch16")
        self.model = AutoModel.from_pretrained("OFA-Sys/chinese-clip-vit-base-patch16")
        
    def get_image_embedding(self, image: Image.Image) -> torch.Tensor:
        inputs = self.processor(images=image, return_tensors="pt")
        with torch.no_grad():
            outputs = self.model.get_image_features(**inputs)
        return outputs.squeeze().numpy() 